from __future__ import annotations
import typing as t

from langchain.chains import LLMChain
from langchain.llms import OpenLLM
from langchain.prompts import PromptTemplate
from pydantic import BaseModel

import bentoml
from bentoml.io import JSON, Text

class Query(BaseModel):
  industry: str
  product_name: str
  keywords: t.List[str]
  llm_config: t.Dict[str, t.Any]

def gen_llm(model_name: str, model_id: str | None = None, **attrs: t.Any) -> OpenLLM:
  lc_llm = OpenLLM(model_name=model_name, model_id=model_id, embedded=False, **attrs)
  lc_llm.runner.download_model()
  return lc_llm

llm = gen_llm("llama", model_id="TheBloke/Llama-2-13B-chat-GPTQ", quantize="gptq")

prompt = PromptTemplate(input_variables=["industry", "product_name", "keywords"],
                        template="""
You are a Facebook Ads Copywriter with a strong background in persuasive
writing and marketing. You craft compelling copy that appeals to the target
audience's emotions and needs, peruading them to take action or make a
purchase. You are given the following context to create a facebook ad copy.
It should provide an attention-grabbing headline optimizied for capivating
leads and perusaive calls to action.

Industry: {industry}
Product: {product_name}
Keywords: {keywords}
Facebook Ads copy:
    """)
chain = LLMChain(llm=llm, prompt=prompt)

svc = bentoml.Service("fb-ads-copy", runners=[llm.runner])

SAMPLE_INPUT = Query(industry="SAAS",
                     product_name="BentoML",
                     keywords=["open source", "developer tool", "AI application platform", "serverless", "cost-efficient"],
                     llm_config=llm.runner.config.model_dump())

@svc.api(input=JSON.from_sample(sample=SAMPLE_INPUT), output=Text())
def generate(query: Query):
  return chain.run({"industry": query.industry, "product_name": query.product_name, "keywords": ", ".join(query.keywords)})
